import dotenv
from langchain_community.tools import GoogleSerperRun
from langchain_community.tools.openai_dalle_image_generation import OpenAIDALLEImageGenerationTool
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
from langchain_core.messages import ToolMessage
from langchain_openai import ChatOpenAI
from langgraph.constants import START
from langgraph.graph import add_messages, StateGraph
from pydantic import BaseModel, Field
import json
from typing import TypedDict, Annotated, Literal

dotenv.load_dotenv()

class GoogleSearchSchema(BaseModel):
    query:str = Field(description="谷歌搜索的参数")

class DalleArgsSchema(BaseModel):
    query: str = Field(description="图片生成指令")

google_search_tool = GoogleSerperRun(
    name= "google_search_tool",
    description= "一个低成本的谷歌搜索API",
    api_wrapper=GoogleSerperAPIWrapper(),
    args_schema=GoogleSearchSchema
)

dalle_tool = OpenAIDALLEImageGenerationTool(
    name= "dalle_tool",
    api_wrapper=DallEAPIWrapper(model="dall-e-3"),
    args_schema=DalleArgsSchema
)

class ChatBotState(TypedDict):
    """聊天消息就是状态，消息是一个字典，add_messages就是一个归纳函数"""
    messages:Annotated[list,add_messages]
bot_tools = [google_search_tool,dalle_tool]
llm =ChatOpenAI(model="gpt-4o-mini")
llm_with_tools = llm.bind_tools(bot_tools)

def chat_bot(state:ChatBotState)->ChatBotState:
    """聊天机器人节点函数"""
    ai_message = llm_with_tools.invoke(state["messages"])

    return {"messages":[ai_message],"node_name":"chat_bot"}

def tool_exe(state:ChatBotState)->ChatBotState:
    """工具执行节点函数"""
    tool_by_name = {tool.name:tool for tool in bot_tools}
    tool_calls = state["messages"][-1].tool_calls


    messages = []
    for tool_call in tool_calls:
        tool = tool_by_name[tool_call["name"]]
        messages.append(ToolMessage(
            tool_call_id = tool_call["id"],
            content = json.dumps(tool.invoke(tool_call["args"])),
            name = tool_call["name"]
        ))
    return {"messages":messages}
def route(state:ChatBotState)->Literal["tool_exe","__end__"]:
    """结束边"""
    ai_message = state["messages"][-1]
    if hasattr(ai_message,"tool_calls") and len(ai_message.tool_calls) >0:
        return "tool_exe"
    return "__end__"
# 构建图架构
graph_builder = StateGraph(ChatBotState)

# 添加聊天节点
graph_builder.add_node("llm",chat_bot)
graph_builder.add_node("tool_exe",tool_exe)
# 添加边
graph_builder.add_edge(START,"llm")
graph_builder.add_conditional_edges("llm",route)
graph_builder.add_edge("tool_exe","llm")
# 编译图架构
graph_app = graph_builder.compile()
# 运行图架构
state = graph_app.invoke({"messages":[("human","帮我生成一张小孩嬉戏图片")]})

for message in state["messages"]:
    print(f"消息类型：{message.type}")
    if hasattr(message,"tool_calls") and len(message.tool_calls) >0:
        print(f"工具调用参数：{message.tool_calls}")
    print(f"消息内容：{message.content}")
    print("================")
